#!/usr/bin/python
'''
Pull matched reads from assembly.
'''
import sys
# parameters.
read_list_file = sys.argv[1]
read_stats_file = sys.argv[2]
reads_file = sys.argv[3]
contig_file = sys.argv[4]
out_dir = sys.argv[5]

def read_file(filepath):
	# define data structure.
	data = {}
	
	# Check if zipped.
	zipped = False
	if filepath.count(".gz"):
		zipped = True
	
	# Open file.
	seq = ""
	seqheader = ""
	first = True
	
	if zipped == True:
		fin = gzip.open(filepath,"r")
	else:
		fin = open(filepath,"r")
	for line in fin:
		# Skip comments.	
		if line[0] == "#": continue

		# first time.
		if first == True and line[0] == ">":
			seq = ""
			seqheader = line.strip().replace(">","")
			first = False
			continue		

		# finished seq.
		if line[0] == ">":
			# save data.
			data[seqheader] = seq.strip()

			# Reset variables.
			seq = ""
			seqheader = line.strip().replace(">","").split(" ")[0]
			continue

		# get seq.
		seq += line.strip()
			
	# Handle last close and return.
	data[seqheader] = seq.strip()
	fin.close()
	return data

# Index reads and contrigs.
reads = read_file(reads_file)
contigs = read_file(contig_file)

# Open read list.
read_list = {}
fin = open(read_list_file,"r")
for line in fin:
	tmp = line.strip().split()
	read_list[tmp[0]] = True	
fin.close()

# Look for reads in assembly.
byctg = {}
fin = open(read_stats_file, "r")
for line in fin:
	tmp = line.strip().split()
	if tmp[0] in read_list:
		if line.count("Assembl") == 0: continue
		
		if tmp[2] not in byctg:
			byctg[tmp[2]] = []
		byctg[tmp[2]].append(tmp[0])
		
# Print only > 5.
for ctg in byctg:
	if len(byctg[ctg]) <= 5: continue
	if ctg not in contigs: continue
	
	fpath = "%s/%s.fasta" % (out_dir, ctg)
	fout = open(fpath,"w")
	
	# write contig.
	fout.write(">%s\n%s\n" % (ctg, contigs[ctg]))
	
	# write reads.
	for read in byctg[ctg]:
		fout.write(">%s\n%s\n" % (read, reads[read]))
	fout.close()
